from __future__ import annotations
from typing import Iterable
import numpy as np

class AudioCPFSK:
    def __init__(self, f_sample: int, f_space: int, f_mark: int, bit_rate: int):
        assert f_sample % bit_rate == 0, "采样频率必须是比特率的整数倍"
        assert (f_space + f_mark) % 2 == 0
        assert 0 < f_space < f_mark < f_sample // 2
        self.f_sample = f_sample
        self.f_space = f_space
        self.f_mark = f_mark
        self.f_center = (f_space + f_mark) / 2
        self.f_delta = (f_mark - f_space) / 2
        self.bit_rate = bit_rate
        self.n_per_bit = int(f_sample // bit_rate)  # 每个比特的采样点数

    def __call__(self, bits: Iterable[bool], amplitude: float=1.0, attack_bits: int=4) -> np.ndarray:
        """
        对比特序列进行CPFSK调制
        Args:
            bits: 比特序列 (True表示1/mark频率, False表示0/space频率)
            amplitude: 信号幅度
            attack_bits: 前导载波的比特数
        Returns:
            调制后的音频信号
        """
        bits_list = list(bits)
        if not bits_list:
            return np.array([])
        
        # 计算前导载波的采样点数
        preamble_samples = attack_bits * self.n_per_bit
        
        # 总采样点数（包括前导）
        data_samples = len(bits_list) * self.n_per_bit
        total_samples = preamble_samples + data_samples
        signal = np.zeros(total_samples)
        
        # 相位累积器，确保相位连续
        phase_accumulator = 0.0
        
        # 1. 生成前导载波信号（使用中心频率）
        preamble_phase_increment = 2 * np.pi * self.f_center / self.f_sample
        for i in range(preamble_samples):
            signal[i] = amplitude * np.cos(phase_accumulator)
            phase_accumulator += preamble_phase_increment
        
        # 保持相位连续性
        phase_accumulator = phase_accumulator % (2 * np.pi)
        
        # 2. 生成数据信号
        for i, bit in enumerate(bits_list):
            # 根据比特值选择频率
            freq = self.f_mark if bit else self.f_space
            
            # 计算当前比特周期内的相位增量
            phase_increment = 2 * np.pi * freq / self.f_sample
            
            # 当前比特的采样点范围（考虑前导偏移）
            start_idx = preamble_samples + i * self.n_per_bit
            
            # 生成当前比特周期的信号
            for j in range(self.n_per_bit):
                signal[start_idx + j] = amplitude * np.cos(phase_accumulator)
                phase_accumulator += phase_increment
                
            # 保持相位在合理范围内，避免数值溢出
            phase_accumulator = phase_accumulator % (2 * np.pi)
        
        return signal
